import wfdb
import numpy as np
import pandas as pd
from scipy.signal import correlate, find_peaks

def load_and_save_mimic_to_csv(record_name, output_csv='processed_data.csv'):
    """加载MIMIC数据并保存为CSV（固定信号名称：PPG=PLETH, ECG=II）"""
    # 读取记录
    record = wfdb.rdrecord(record_name)
    
    # 硬编码信号名称（根据您的确认）
    PPG_NAME = 'PLETH'
    ECG_NAME = 'II'
    ABP_NAME = 'ABP'
    
    # 验证信号存在性
    assert PPG_NAME in record.sig_name, f"Required signal '{PPG_NAME}' not found. Available: {record.sig_name}"
    assert ABP_NAME in record.sig_name, f"Required signal '{ABP_NAME}' not found. Available: {record.sig_name}"
    assert ECG_NAME in record.sig_name, f"Required signal '{ECG_NAME}' not found. Available: {record.sig_name}"

    # 提取信号
    signals = {
        'ppg': record.p_signal[:, record.sig_name.index(PPG_NAME)],
        'abp': record.p_signal[:, record.sig_name.index(ABP_NAME)],
        'ecg': record.p_signal[:, record.sig_name.index(ECG_NAME)]
    }

    for key in signals:
        sig = signals[key]
        if np.isnan(sig).any() or np.isinf(sig).any():
            print(f"信号 {key} 含有 NaN 或 Inf，正在进行修复...")
            mean_val = np.nanmean(sig[np.isfinite(sig)])  # 仅对有效值取均值
            signals[key] = np.nan_to_num(sig, nan=mean_val, posinf=mean_val, neginf=mean_val)

    time = np.arange(len(signals['ppg'])) / record.fs  # 时间轴（秒）


    # PPG-ABP对齐
    cross_corr = correlate(signals['ppg'] - signals['ppg'].mean(),
                          signals['abp'] - signals['abp'].mean(),
                          mode='full')
    delay = cross_corr.argmax() - (len(signals['ppg']) - 1)
    abp_aligned = np.roll(signals['abp'], -delay)

    # 使用ECG (II导联)检测R峰
    ecg_peaks, _ = find_peaks(signals['ecg'], 
                            height=np.nanpercentile(signals['ecg'],80),
                            distance=record.fs*0.6) # 最小间隔0.6秒
    delta = np.mean(np.diff(ecg_peaks)) / record.fs if len(ecg_peaks) > 1 else 0.8  # 默认周期0.8秒
    
    # 构建DataFrame
    df = pd.DataFrame({
        'time': time,
        'pleth': signals['ppg'],       # PLETH信号
        'abp_raw': signals['abp'],     # 原始ABP
        'abp_aligned': abp_aligned,    # 对齐后的ABP
        'ecg_ii': signals['ecg'],      # II导联ECG
        'r_peak': np.isin(np.arange(len(time)), ecg_peaks)  # R峰标记
    })

    # 保存CSV（含元数据注释）
    meta = (f"# fs={record.fs}, delay_samples={delay}, "
            f"ppg_name={PPG_NAME}, ecg_name={ECG_NAME}, "
            f"hr={60/delta:.1f}bpm\n")
    with open(output_csv, 'w') as f:
        f.write(meta)
        df.to_csv(f, index=False)

    print(f"数据已保存至 {output_csv}")
    print(f"信号长度: {len(signals['ppg'])/record.fs:.1f}秒 | 心率: {60/delta:.1f}bpm")

# 示例使用
load_and_save_mimic_to_csv('3300784_0003', 'original_data_0003.csv')
#你可以把你上面提出的这三个建议融进这个代码嘛